#include "hr.hpp"
#ifndef __SYNTHESIS__ // 仅在 C 仿真时包含 iostream
#include <iostream>
#endif

// 辅助函数：判断点是否在三角形内（使用边缘方程）
bool is_inside_triangle(
    fpal::fp_coord_t x, fpal::fp_coord_t y,
    fpal::fp_coord_t A0, fpal::fp_coord_t B0, fpal::fp_coord_t C0,
    fpal::fp_coord_t A1, fpal::fp_coord_t B1, fpal::fp_coord_t C1,
    fpal::fp_coord_t A2, fpal::fp_coord_t B2, fpal::fp_coord_t C2
) {
    #pragma HLS INLINE 
    fpal::fp_coord_t w0 = fpal::add_fx(fpal::mul_fx(A0, x), fpal::add_fx(fpal::mul_fx(B0, y), C0));
    fpal::fp_coord_t w1 = fpal::add_fx(fpal::mul_fx(A1, x), fpal::add_fx(fpal::mul_fx(B1, y), C1));
    fpal::fp_coord_t w2 = fpal::add_fx(fpal::mul_fx(A2, x), fpal::add_fx(fpal::mul_fx(B2, y), C2));

    bool inside = (w0 >= 0 && w1 >= 0 && w2 >= 0);

    return inside;
}

// 辅助函数：更新Bounding Box
void update_bounding_box(
    fpal::fp_coord_t vx, fpal::fp_coord_t vy,
    fpal::fp_screen_pixel_idx_t& min_x, fpal::fp_screen_pixel_idx_t& max_x,
    fpal::fp_screen_pixel_idx_t& min_y, fpal::fp_screen_pixel_idx_t& max_y,
    const ap_uint<11> max_x_res, const ap_uint<11> max_y_res
) {
    #pragma HLS INLINE
    fpal::fp_coord_t max_x_res_fx = fpal::to_fp_coord(max_x_res);
    fpal::fp_coord_t max_y_res_fx = fpal::to_fp_coord(max_y_res);
    fpal::fp_coord_t one_fx = fpal::to_fp_coord(1); 
    
    fpal::fp_coord_t clamped_vx = fpal::clamp(vx, fpal::to_fp_coord(0), fpal::sub_fx(max_x_res_fx, one_fx));
    fpal::fp_coord_t clamped_vy = fpal::clamp(vy, fpal::to_fp_coord(0), fpal::sub_fx(max_y_res_fx, one_fx));

    fpal::fp_screen_pixel_idx_t cur_x = fpal::to_fp_screen_pixel_idx(clamped_vx);
    fpal::fp_screen_pixel_idx_t cur_y = fpal::to_fp_screen_pixel_idx(clamped_vy);

    if (cur_x < min_x) min_x = cur_x;
    if (cur_x > max_x) max_x = cur_x;
    if (cur_y < min_y) min_y = cur_y;
    if (cur_y > max_y) max_y = cur_y;
}


// ===========================================================
// HR_Setup 模块 (Sub-process 1)
// 职责：从输入流读取一个三角形的参数，计算其包围盒、边缘方程系数，
//       深度梯度，并将这些“三角形级”的参数打包成 triangle_setup_packet_fixed 推送到中间流。
// ===========================================================
void hr_setup(
    hls::stream<raster_params_packet_fixed>& input_raster_params_stream,
    hls::stream<triangle_setup_packet_fixed>& output_triangle_setup_stream,
    const ap_uint<11> max_x_res, 
    const ap_uint<11> max_y_res,  
    const ap_uint<16> num_triangles_to_process
) {
    #pragma HLS DATAFLOW // This sub-function should be part of a larger DATAFLOW region

    TRIANGLE_SETUP_LOOP:
    for (ap_uint<16> triangle_idx = 0; triangle_idx < num_triangles_to_process; ++triangle_idx) {
        #pragma HLS LOOP_TRIPCOUNT min=500 max=2000
        #pragma HLS PIPELINE II=1

        raster_params_packet_fixed current_triangle_params;
        input_raster_params_stream.read(current_triangle_params);

#ifndef __SYNTHESIS__
        std::cout << "[HR_SETUP_DEBUG] Processing Triangle " << triangle_idx << std::endl;
#endif

        fpal::fp_coord_t v0_x_fx = current_triangle_params.v0_x;
        fpal::fp_coord_t v0_y_fx = current_triangle_params.v0_y;

        fpal::fp_coord_t max_x_res_fx = fpal::to_fp_coord(max_x_res);
        fpal::fp_coord_t max_y_res_fx = fpal::to_fp_coord(max_y_res);
        fpal::fp_coord_t one_fx = fpal::to_fp_coord(1);
        fpal::fp_coord_t zero_fx = fpal::to_fp_coord(0);

        fpal::fp_screen_pixel_idx_t min_x = fpal::to_fp_screen_pixel_idx(fpal::clamp(v0_x_fx, zero_fx, fpal::sub_fx(max_x_res_fx, one_fx)));
        fpal::fp_screen_pixel_idx_t max_x = min_x;
        fpal::fp_screen_pixel_idx_t min_y = fpal::to_fp_screen_pixel_idx(fpal::clamp(v0_y_fx, zero_fx, fpal::sub_fx(max_y_res_fx, one_fx)));
        fpal::fp_screen_pixel_idx_t max_y = min_y;

        update_bounding_box(current_triangle_params.v1_x, current_triangle_params.v1_y, 
                            min_x, max_x, min_y, max_y, max_x_res, max_y_res);
        update_bounding_box(current_triangle_params.v2_x, current_triangle_params.v2_y, 
                            min_x, max_x, min_y, max_y, max_x_res, max_y_res);
        
        min_x = fpal::clamp(min_x, (fpal::fp_screen_pixel_idx_t)0, fpal::to_fp_screen_pixel_idx(fpal::sub_fx(max_x_res_fx, one_fx)));
        min_y = fpal::clamp(min_y, (fpal::fp_screen_pixel_idx_t)0, fpal::to_fp_screen_pixel_idx(fpal::sub_fx(max_y_res_fx, one_fx)));
        max_x = fpal::clamp(max_x, (fpal::fp_screen_pixel_idx_t)0, fpal::to_fp_screen_pixel_idx(fpal::sub_fx(max_x_res_fx, one_fx)));
        max_y = fpal::clamp(max_y, (fpal::fp_screen_pixel_idx_t)0, fpal::to_fp_screen_pixel_idx(fpal::sub_fx(max_y_res_fx, one_fx)));
        
        triangle_setup_packet_fixed setup_packet;
        setup_packet.v0_x = v0_x_fx;
        setup_packet.v0_y = v0_y_fx;
        setup_packet.v0_z = current_triangle_params.v0_z;
        setup_packet.edge0_A = current_triangle_params.edge0_A; setup_packet.edge0_B = current_triangle_params.edge0_B; setup_packet.edge0_C = current_triangle_params.edge0_C;
        setup_packet.edge1_A = current_triangle_params.edge1_A; setup_packet.edge1_B = current_triangle_params.edge1_B; setup_packet.edge1_C = current_triangle_params.edge1_C;
        setup_packet.edge2_A = current_triangle_params.edge2_A; setup_packet.edge2_B = current_triangle_params.edge2_B; setup_packet.edge2_C = current_triangle_params.edge2_C;
        setup_packet.dzdx = current_triangle_params.dzdx;
        setup_packet.dzdy = current_triangle_params.dzdy;
        setup_packet.pure_color_R = current_triangle_params.pure_color_R;
        setup_packet.pure_color_G = current_triangle_params.pure_color_G;
        setup_packet.pure_color_B = current_triangle_params.pure_color_B;
        setup_packet.min_x = min_x; setup_packet.max_x = max_x;
        setup_packet.min_y = min_y; setup_packet.max_y = max_y;
        setup_packet.max_x_res = max_x_res; // Pass max_x_res for hr_pixel_generator

        output_triangle_setup_stream.write(setup_packet);

#ifndef __SYNTHESIS__
        std::cout << "[HR_SETUP_DEBUG]   Bounding Box: (" << min_x << "," << min_y << ") to (" << max_x << "," << max_y << ")" << std::endl;
#endif
    }
}


// ===========================================================
// HR_Row_Processor 模块 (Sub-process 2)
// 职责：接收上一个模块的三角形参数。它将负责 Y 轴上的循环。
//       对于包围盒内的每一行，它会根据深度梯度计算该行的起始像素深度，
//       并打包成一个“行描述”结构体，推送到中间流。
// ===========================================================
void hr_row_processor(
    hls::stream<triangle_setup_packet_fixed>& input_triangle_setup_stream,
    hls::stream<row_description_packet_fixed>& output_row_description_stream,
    const ap_uint<16> num_triangles_to_process,
    hls::stream<ap_uint<32>>& output_estimated_fragments_count
) {
    #pragma HLS DATAFLOW

    ap_uint<32> fragment_counter_local = 0; // 局部计数器

    ROW_PROCESS_TRIANGLE_LOOP:
    for (ap_uint<16> triangle_idx = 0; triangle_idx < num_triangles_to_process; ++triangle_idx) {
        #pragma HLS LOOP_TRIPCOUNT min=500 max=2000
        #pragma HLS PIPELINE II=1 // Loop for processing each triangle's rows

        triangle_setup_packet_fixed setup_packet;
        input_triangle_setup_stream.read(setup_packet);

#ifndef __SYNTHESIS__
        std::cout << "[HR_ROW_PROC_DEBUG] Processing rows for Triangle " << triangle_idx 
                  << " (BB: Y from " << setup_packet.min_y << " to " << setup_packet.max_y << ")" << std::endl;
#endif

        fpal::fp_coord_t v0_x_fx = setup_packet.v0_x;
        fpal::fp_coord_t v0_y_fx = setup_packet.v0_y;
        fpal::fp_coord_t v0_z_fx = fpal::to_fp_coord(setup_packet.v0_z); // Convert V0's Z to fp_coord_t for interpolation

        // Calculate initial Z at (min_x, min_y) for the triangle
        fpal::fp_coord_t min_x_fx = fpal::to_fp_coord(setup_packet.min_x);
        fpal::fp_coord_t min_y_fx = fpal::to_fp_coord(setup_packet.min_y);

        fpal::fp_coord_t delta_x_initial = fpal::sub_fx(min_x_fx, v0_x_fx);
        fpal::fp_coord_t delta_y_initial = fpal::sub_fx(min_y_fx, v0_y_fx);
        
        fpal::fp_coord_t initial_z_offset_x = fpal::mul_fx(setup_packet.dzdx, delta_x_initial);
        fpal::fp_coord_t initial_z_offset_y = fpal::mul_fx(setup_packet.dzdy, delta_y_initial);
        fpal::fp_coord_t initial_z_at_bb_min_xy = fpal::add_fx(v0_z_fx, fpal::add_fx(initial_z_offset_x, initial_z_offset_y));

        fpal::fp_coord_t current_z_at_row_start_fx = initial_z_at_bb_min_xy;
        
        // Accumulate dzdy across the rows to get initial_z_for_row_fx efficiently
        Y_LOOP: 
        for (fpal::fp_screen_pixel_idx_t y = setup_packet.min_y; y <= setup_packet.max_y; ++y) {
            #pragma HLS LOOP_TRIPCOUNT min=1 max=720
            #pragma HLS PIPELINE II=1

            row_description_packet_fixed row_packet;
            row_packet.y = y;
            row_packet.min_x = setup_packet.min_x;
            row_packet.max_x = setup_packet.max_x;

            // Z interpolation for the start of the current row (min_x, y)
            // (initial_z_at_bb_min_xy + (y - min_y) * dzdy)
            fpal::fp_coord_t delta_y_row = fpal::sub_fx(fpal::to_fp_coord(y), min_y_fx);
            row_packet.initial_z_for_row_fx = fpal::add_fx(initial_z_at_bb_min_xy, fpal::mul_fx(delta_y_row, setup_packet.dzdy));
            
            row_packet.dzdx = setup_packet.dzdx; // dzdx is constant for the triangle

            row_packet.pure_color_R = setup_packet.pure_color_R;
            row_packet.pure_color_G = setup_packet.pure_color_G;
            row_packet.pure_color_B = setup_packet.pure_color_B;

            row_packet.edge0_A = setup_packet.edge0_A; row_packet.edge0_B = setup_packet.edge0_B; row_packet.edge0_C = setup_packet.edge0_C;
            row_packet.edge1_A = setup_packet.edge1_A; row_packet.edge1_B = setup_packet.edge1_B; row_packet.edge1_C = setup_packet.edge1_C;
            row_packet.edge2_A = setup_packet.edge2_A; row_packet.edge2_B = setup_packet.edge2_B; row_packet.edge2_C = setup_packet.edge2_C;
            
            output_row_description_stream.write(row_packet);
            
            // Count fragments that *could* be generated for this row
            if (setup_packet.min_x <= setup_packet.max_x) { // Only if bounding box is valid horizontally
                 fragment_counter_local += (setup_packet.max_x - setup_packet.min_x + 1);
            }

#ifndef __SYNTHESIS__
            if (y < setup_packet.min_y + 10 || y == setup_packet.max_y) {
                std::cout << "[HR_ROW_PROC_DEBUG]     Row " << y << ", X from " << row_packet.min_x << " to " << row_packet.max_x 
                          << ", Init Z (fp_coord): " << row_packet.initial_z_for_row_fx << std::endl;
            }
#endif
        }
    }
    // Pass the estimated total fragment count (based on bounding boxes) to the next stage
    output_estimated_fragments_count.write(fragment_counter_local); // MODIFIED: Write to stream
}

// ===========================================================
// HR_Pixel_Generator 模块 (Sub-process 3)
// 职责：接收“行描述”结构体。它将负责 X 轴上的循环。
//       对于行内的每一个像素，它会进行三角形内部测试、深度插值，
//       并生成最终的 pixel_fragment_packet_fixed，推送到 output_pixel_fragment_stream。
// ===========================================================
void hr_pixel_generator(
    hls::stream<row_description_packet_fixed>& input_row_description_stream,
    hls::stream<pixel_fragment_packet_fixed>& output_pixel_fragment_stream,
    hls::stream<ap_uint<32>>& input_fragments_count_hr_row_proc, // 从 hr_row_processor 接收总片段数
    hls::stream<ap_uint<32>>& output_actual_fragments_count_stream // hr_pixel_generator 写入实际片段计数
) {
    #pragma HLS DATAFLOW

    ap_uint<32> num_fragments_to_process_estimate;
    input_fragments_count_hr_row_proc.read(num_fragments_to_process_estimate);
#ifndef __SYNTHESIS__
    std::cout << "[HR_PIXEL_GEN_DEBUG] Pixel generator expects processing up to: " << num_fragments_to_process_estimate << " fragments." << std::endl;
#endif
    
    ap_uint<32> hr_total_fragments_generated_actual = 0; 
    
    PIXEL_GEN_TOTAL_LOOP:
    for (ap_uint<32> row_idx = 0; row_idx < num_fragments_to_process_estimate;) {
        #pragma HLS LOOP_TRIPCOUNT min=100000 max=944754
        #pragma HLS PIPELINE II=1

        row_description_packet_fixed row_packet;
        input_row_description_stream.read(row_packet); 
        
#ifndef __SYNTHESIS__
        if (row_idx < 10 || row_idx % 1000 == 0) {
            std::cout << "[HR_PIXEL_GEN_DEBUG]   Processing Row " << row_packet.y 
                      << ", X from " << row_packet.min_x << " to " << row_packet.max_x << std::endl;
        }
#endif
        
        X_LOOP: 
        for (fpal::fp_screen_pixel_idx_t x = row_packet.min_x; x <= row_packet.max_x; ++x) {
            #pragma HLS LOOP_TRIPCOUNT min=1 max=1280 
            #pragma HLS PIPELINE II=1

            pixel_fragment_packet_fixed fragment;
            fragment.screen_x = x;
            fragment.screen_y = row_packet.y;
            fragment.pure_color_R = row_packet.pure_color_R;
            fragment.pure_color_G = row_packet.pure_color_G;
            fragment.pure_color_B = row_packet.pure_color_B;
            fragment.is_valid_pixel = false; 

            fpal::fp_coord_t pixel_x_coord = fpal::to_fp_coord(x);
            fpal::fp_coord_t pixel_y_coord = fpal::to_fp_coord(row_packet.y);

            if (is_inside_triangle(
                    pixel_x_coord, pixel_y_coord,
                    row_packet.edge0_A, row_packet.edge0_B, row_packet.edge0_C,
                    row_packet.edge1_A, row_packet.edge1_B, row_packet.edge1_C,
                    row_packet.edge2_A, row_packet.edge2_B, row_packet.edge2_C
                )) {
                
                fragment.is_valid_pixel = true;

                // 深度插值 (线性插值)
                // current_z_coord_val_fx = (initial_z_for_row_fx + (x - min_x) * dzdx)
                fpal::fp_coord_t delta_x_pixel = fpal::sub_fx(pixel_x_coord, fpal::to_fp_coord(row_packet.min_x));
                fpal::fp_coord_t dz_offset_x = fpal::mul_fx(row_packet.dzdx, delta_x_pixel);
                fpal::fp_coord_t interpolated_z_fx = fpal::add_fx(row_packet.initial_z_for_row_fx, dz_offset_x);
                
                fragment.interpolated_z = fpal::to_fp_z_buffer(interpolated_z_fx);

#ifndef __SYNTHESIS__
                if (hr_total_fragments_generated_actual < 10 || hr_total_fragments_generated_actual % 10000 == 0) {
                    std::cout << "[HR_PIXEL_GEN_DEBUG]     Valid Pixel (" << x << "," << row_packet.y << ") -> Interp. Z: " 
                              << fragment.interpolated_z << std::endl;
                }
#endif
            }
            output_pixel_fragment_stream.write(fragment);
            hr_total_fragments_generated_actual++; 
        }
        if (row_packet.min_x <= row_packet.max_x) {
             row_idx += (row_packet.max_x - row_packet.min_x + 1);
        } else {
             row_idx += 1;
        }
    }

    output_actual_fragments_count_stream.write(hr_total_fragments_generated_actual); // <-- MODIFIED: 写入实际计数到新的输出流
}


// ===========================================================
// HR 模块的顶层 HLS 函数实现
// ===========================================================
void hardware_rasterizer(
    hls::stream<raster_params_packet_fixed>& input_raster_params_stream,
    hls::stream<pixel_fragment_packet_fixed>& output_pixel_fragment_stream,
    const ap_uint<11> max_x_res, 
    const ap_uint<11> max_y_res,  
    const ap_uint<16> num_triangles_to_process,
    ap_uint<32>& total_fragments_generated,         
    hls::stream<ap_uint<32>>& hr_to_pdtd_fragment_count_stream // PDTD 的输入流，由 HR 顶层写入
) {
    #pragma HLS INTERFACE axis port=input_raster_params_stream
    #pragma HLS INTERFACE axis port=output_pixel_fragment_stream
    #pragma HLS INTERFACE axis port=hr_to_pdtd_fragment_count_stream

    #pragma HLS INTERFACE s_axilite port=max_x_res bundle=ctrl_s_axilite_hr
    #pragma HLS INTERFACE s_axilite port=max_y_res bundle=ctrl_s_axilite_hr
    #pragma HLS INTERFACE s_axilite port=num_triangles_to_process bundle=ctrl_s_axilite_hr
    #pragma HLS INTERFACE s_axilite port=total_fragments_generated bundle=ctrl_s_axilite_hr
    #pragma HLS INTERFACE s_axilite port=return bundle=ctrl_s_axilite_hr
                     
    #pragma HLS DATAFLOW

    // HR 内部的 Stream 声明
    hls::stream<triangle_setup_packet_fixed> hr_setup_to_row_processor_stream("hr_setup_to_row_processor_stream");
    #pragma HLS STREAM variable=hr_setup_to_row_processor_stream depth=2 

    hls::stream<row_description_packet_fixed> hr_row_processor_to_pixel_generator_stream("hr_row_processor_to_pixel_generator_stream");
    #pragma HLS STREAM variable=hr_row_processor_to_pixel_generator_stream depth=64 

    // 用于在 hr_row_processor 和 hr_pixel_generator 之间传递总片段的估计值
    hls::stream<ap_uint<32>> hr_row_proc_to_pixel_gen_estimate_stream("hr_row_proc_to_pixel_gen_estimate_stream"); // <-- 重命名
    #pragma HLS STREAM variable=hr_row_proc_to_pixel_gen_estimate_stream depth=2

    // 用于从 hr_pixel_generator 接收实际生成的片段数
    hls::stream<ap_uint<32>> hr_pixel_gen_to_hr_top_actual_count_stream("hr_pixel_gen_to_hr_top_actual_count_stream"); // <-- 新增流
    #pragma HLS STREAM variable=hr_pixel_gen_to_hr_top_actual_count_stream depth=2


    // 实例化并连接子模块
    hr_setup(
        input_raster_params_stream,
        hr_setup_to_row_processor_stream,
        max_x_res,
        max_y_res,
        num_triangles_to_process
    );

    // hr_row_processor 直接写入流
    hr_row_processor(
        hr_setup_to_row_processor_stream,
        hr_row_processor_to_pixel_generator_stream,
        num_triangles_to_process,
        hr_row_proc_to_pixel_gen_estimate_stream // hr_row_processor 直接写入此流
    );

    hr_pixel_generator(
        hr_row_processor_to_pixel_generator_stream,
        output_pixel_fragment_stream,
        hr_row_proc_to_pixel_gen_estimate_stream, 
        hr_pixel_gen_to_hr_top_actual_count_stream // hr_pixel_generator 写入此流
    );

    // 从 hr_pixel_generator 接收实际片段计数
    ap_uint<32> final_actual_fragment_count_from_pixel_gen; // 局部变量
    hr_pixel_gen_to_hr_top_actual_count_stream.read(final_actual_fragment_count_from_pixel_gen); // 从新流读取
    total_fragments_generated = final_actual_fragment_count_from_pixel_gen; // 赋给顶层输出参数
    hr_to_pdtd_fragment_count_stream.write(final_actual_fragment_count_from_pixel_gen); // 将计数传递给 PDTD

#ifndef __SYNTHESIS__
    std::cout << "[HR DEBUG] Total fragments generated by HR (final reported): " << total_fragments_generated << std::endl;
#endif
}
